/**
 * Copyright 2019-2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef FUSION_ENGINE_OPTIMIZER_GRAPH_OPTIMIZER_SHAPE_FORMAT_TRANSFER_TRANS_NODE_IMPLEMENTATION_TRANS_NODE_BASE_GENERATOR_H_
#define FUSION_ENGINE_OPTIMIZER_GRAPH_OPTIMIZER_SHAPE_FORMAT_TRANSFER_TRANS_NODE_IMPLEMENTATION_TRANS_NODE_BASE_GENERATOR_H_
#include "common/fe_inner_attr_define.h"
#include "common/fe_inner_error_codes.h"
#include "graph/compute_graph.h"
#include "graph/debug/ge_attr_define.h"
#include "graph/utils/graph_utils.h"
#include "ops_kernel_store/fe_ops_kernel_info_store.h"

namespace fe {
struct TransferInfo {
  ge::OpDescPtr src_op_desc;
  ge::OutDataAnchorPtr src_anchor;
  ge::OpDescPtr dst_op_desc;
  ge::InDataAnchorPtr dst_anchor;
  ge::NodePtr src_node_ptr;
  ge::NodePtr dst_node_ptr;
  ge::ConstGeTensorDescPtr src_out_tensor_desc_ptr;
  ge::ConstGeTensorDescPtr dst_in_tensor_desc_ptr;

  ge::Format src_out_primary_format;
  int32_t src_out_sub_format = 0;
  ge::DataType src_out_data_type;
  ge::GeShape src_out_shape;

  ge::Format dst_in_primary_format;
  int32_t dst_in_sub_format = 0;
  ge::DataType dst_in_data_type;
  ge::GeShape dst_in_shape;

  string src_op_desc_type;
  string dst_op_desc_type;

  string src_reshape_type;
  string dst_reshape_type;


  /* src_imply_type is always pointed to the source normal op in the origin graph
   * instead of trans op which is inserted by this class TransOpInsert */
  OpImplType src_imply_type;
  /* dst_imply_type is always pointed to the destination normal op in the origin
   * graph instead of trans op which is inserted by this class TransOpInsert */
  OpImplType dst_imply_type;

  uint32_t insertion_mode;

  /* Set this source out orginal format and shape for getting 4D shape
   * without padding in 5->4 scenario */
  ge::Format src_out_original_format;
  ge::GeShape src_out_original_shape;

  ge::Format dst_in_original_format;
  ge::GeShape dst_in_original_shape;

  bool is_source_weight;
  bool is_dst_weight;

  OpPattern src_op_pattern;
  OpPattern dst_op_pattern;

  vector<std::pair<int64_t, int64_t>> src_out_range;
  vector<std::pair<int64_t, int64_t>> dst_in_range;
};

using FEOpsKernelInfoStorePtr = std::shared_ptr<FEOpsKernelInfoStore>;
using TransInfo = struct TransferInfo;
using TransInfoPtr = std::shared_ptr<TransInfo>;

/** @brief Trans-node base class. Provide function of setting basic op_desc and
* attributes and adding nodes nad edges. It will inherited by specific
* trans-nodes like Cast, Reshape, TransData, Transpose. */
class TransNodeBaseGenerator {
 public:
  TransNodeBaseGenerator(FEOpsKernelInfoStorePtr fe_ops_store_ptr, TransInfoPtr trans_info_ptr);

  virtual ~TransNodeBaseGenerator();

  TransNodeBaseGenerator(const TransNodeBaseGenerator &) = delete;

  TransNodeBaseGenerator &operator=(const TransNodeBaseGenerator &) = delete;

  /* Create a Opdesc Pointer for All trans-ops. It includes all common info
   * like of input, output, fe imply type, ge imply type e.g. of trans-nodes
   * such as cast, reshape, transpose, transdata.
   * */
  ge::OpDescPtr CreateBasicOpDescForTransNode(const string &op_type);

  Status AddEdgesAndFreshTransInfo(ge::ComputeGraph &fused_graph, ge::OpDescPtr &op_desc);

  /* Add specific op info into OpDescPtr. It's implemented by different sub
   * classes because different trans-nodes have different attributes. */
  virtual Status AddTransNode(ge::ComputeGraph &fused_graph, TransInfoPtr trans_info_ptr) = 0;

  Status TransformDimTo4(bool increasing_flag);

  virtual Status SetTensorDescInfo(ge::OpDescPtr &op_desc_ptr);

  Status SetTensorRealDimCountAndNewShape(ge::OpDescPtr &op_desc_ptr, std::vector<ge::GeShape> inputs_shape,
                                          ge::GeShape output_shape);

  Status SetNewShapeRange(ge::OpDescPtr &op_desc_ptr, vector<std::pair<int64_t, int64_t>> &inputs_range,
                          vector<std::pair<int64_t, int64_t>> &output_range);

  TransInfoPtr trans_info_ptr_;

  bool TransNodeCheckAccuracySupported(ge::OpDescPtr &op_desc_ptr, bool real_query,
                                       bool not_need_check_support_flag = false);

  static uint64_t GetTransAtomicId();
 private:
  /* Add necessory peer nodes */
  virtual Status AddNecessaryPeerNodes(ge::ComputeGraph &fused_graph, ge::NodePtr new_node);
  /* Add this trans-node into graph and add its input and output edges. */
  Status AddEdgesForNewNode(ge::NodePtr new_node);

  void RefreshSourceTransInfo(ge::NodePtr src_node);

  FEOpsKernelInfoStorePtr fe_ops_store_info_ptr_;
};

}  // namespace fe
#endif  // FUSION_ENGINE_OPTIMIZER_GRAPH_OPTIMIZER_SHAPE_FORMAT_TRANSFER_TRANS_NODE_IMPLEMENTATION_TRANS_NODE_BASE_GENERATOR_H_